### Analytical solutions and approximations ####################################

from numpy import power,sqrt

################################## Stability of CI #############################

### Convenience functions ######################################################
def R(f,ci,t):
    """Convenience function R.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return power(ci-f,2) - 4 * power(1-f,2) * ci * t * (1-t)

def Rh(f,ci,m,t):
    """Convenience function R^.
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    return power(ci-f,2) - 4 * (1-f) * ci * t * ( m + (1-f)*(1-m)*(1-t) )

def D(f,ci,t):
    """Convenience function D.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return 2 * ci * ( f + (1-f)*t )

def P(f,ci,m,t):
    """Convenience function.
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    return (f+ci) * (1+m) - (1-m) * sqrt(R(f,ci,t))

def Q(f,t):
    """Convenience function Q.
    f  - fecundity reduction in infected females
    t  - transmission rate of Wolbachia"""
    return 2 * power(1-f,2) * t * (1-t)

def Qh(f,m,t):
    """Convenience function Q^.
    f  - fecundity reduction in infected females
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    return 2 * (1-f) * ( (1 - f*(1-m)) * t - (1-f)*(1-m)*power(t,2) )

### Single host population #####################################################
# Dynamics:
def F(f,ci,t,x):
    """Infection dynamics F of Wolbachia, iterative map returns the infection 
    frequency in the next generation.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia
    x  - infection frequency of Wolbachia"""
    return (1-f)*t*x / (1-f*x*(1-ci*(1-t)*x)-ci*x*(1-t*x))
    
# Fixpoints:
def fix1SP(f,ci,t):
    """Fixpoint x1* for a single host population (SP).
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return 0.

def fix2SP(f,ci,t):
    """Fixpoint x2* for a single host population (SP).
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    numer = f + ci - sqrt(R(f,ci,t))
    denom = D(f,ci,t)
    return numer/denom

def fix3SP(f,ci,t):
    """Fixpoint x3* for a single host population (SP).
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    numer = f + ci + sqrt(R(f,ci,t))
    denom = D(f,ci,t)
    return numer/denom

# Critical CI level:
def lcSP(f,t):
    """Critical CI level (lc) for a single host population (SP).
    f  - fecundity reduction in infected females
    t  - transmission rate of Wolbachia"""
    return f + Q(f,t) + sqrt( power(f+Q(f,t),2) - power(f,2) )

### Uninfected mainland ########################################################
# Dynamics:
def FUM(f,ci,m,t,x):
    """Infection dynamics G of Wolbachia for the scenario with an uninfected 
    mainland (UM), iterative map returns the infection frequency in the next 
    generation.
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia
    x  - infection frequency of Wolbachia"""
    return (1-m)*F(f,ci,t,x)
    
# Fixpoints:
def fix1UM(f,ci,m,t):
    """Fixpoint x1* for the scenario with an uninfected mainland (UM).
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    return 0.

def fix2UM(f,ci,m,t):
    """Fixpoint x2* for the scenario with an uninfected mainland (UM).
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    numer = f + ci - sqrt(Rh(f,ci,m,t))
    denom = D(f,ci,t)
    return numer/denom

def fix3UM(f,ci,m,t):
    """Fixpoint x3* for the scenario with an uninfected mainland (UM).
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    numer = f + ci + sqrt(Rh(f,ci,m,t))
    denom = D(f,ci,t)
    return numer/denom

# Critical CI level:
def lcUM(f,m,t):
    """Critical CI level (lc) for the scenario with an uninfected mainland (UM).
    f  - fecundity reduction in infected females
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    return f + Qh(f,m,t) + sqrt( power(f+Qh(f,m,t),2) - power(f,2) )

# Critical migration rate:
def mcUM(f,ci,t):
    """Critical migration rate (mc) for the scenario with an uninfected 
    mainland (UM).
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    numer = R(f,ci,t)
    denom = 2 * (1-f) * t * D(f,ci,t)
    return numer/denom

### Infected mainland ##########################################################
# Dynamics:
def FIM(f,ci,m,t,x):
    """Infection dynamics H of Wolbachia for the scenario with an infected 
    mainland (IM), iterative map.
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia
    x  - infection frequency of Wolbachia"""
    return (1-m)*F(f,ci,t,x) + m*fix3SP(f,ci,t)
    
# Fixpoints:
def fix1IM(f,ci,m,t):
    """Fixpoint x1* for the scenario with an infected mainland (IM).
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    return fix3SP(f,ci,t)

def fix2IM(f,ci,m,t):
    """Fixpoint x2* for the scenario with an infected mainland (IM).
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    numer = P(f,ci,m,t) + sqrt( power(P(f,ci,m,t),2) - 8 * m * D(f,ci,t) )
    denom = 2*D(f,ci,t)
    return numer/denom

def fix3IM(f,ci,m,t):
    """Fixpoint x3* for the scenario with an infected mainland (IM).
    f  - fecundity reduction in infected females
    ci - level of CI
    m  - migration rate
    t  - transmission rate of Wolbachia"""
    numer = P(f,ci,m,t) - sqrt( power(P(f,ci,m,t),2) - 8 * m * D(f,ci,t) )
    denom = 2*D(f,ci,t)
    return numer/denom

# Critical migration rate:
def mcIM(f,ci,t):
    """Critical migration rate (mc) for the scenario with an infected mainland 
    (IM).
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    numer = 1 - sqrt((1-f)*t)
    denom = f + ci + sqrt(R(f,ci,t))
    return 2 * D(f,ci,t) * power( numer/denom, 2 )

### Local host adaptation ######################################################
# Critical migration rates:
def mcUMA(f,ci,s,t):
    """Critical migration rate (mc) for the scenario with an uninfected 
    mainland (UM) with local host adaptation (A).
    f  - fecundity reduction in infected females
    ci - level of CI
    s  - selection coefficent
    t  - transmission rate of Wolbachia"""
    numer = (1+s) * R(f,ci,t)
    denom = 2 * (1-f) * t * D(f,ci,t)  +  s * R(f,ci,t)
    return numer/denom

def mcIMA(f,ci,s,t):
    """Critical migration rate (mc) for the scenario with an infected mainland 
    (IM) with local host adaptation (A).
    f  - fecundity reduction in infected females
    ci - level of CI
    s  - selection coefficent
    t  - transmission rate of Wolbachia"""
    numer = (1+s) * ( 2 * D(f,ci,t) * power(1-sqrt((1-f)*t),2) + s * power(f+ci-sqrt(R(f,ci,t)),2) )
    denom = power( (f+ci)*(1-s) + (1+s)*sqrt(R(f,ci,t)), 2 ) + 8 * s * D(f,ci,t)
    return  numer/denom


################################## Gene flow ###################################
# Divergent selection:
def rvDS(s):
    """Reproductive value (rv) of a migrant for the case of divergent selection 
    (DS). Derived from fitness graph.
    s - selection coefficient: residents have a viability advantage of s over 
        migrants, equivalent to migrants having a viability cost of 
        s^ = s/(1+s)"""
    return 1./(1+2*s)

def meDS(m,s):
    """Effective migration rate (me) for the case of divergent selection (DS).
    Derived from fitness graph for the reproductive value of a migrant.
    m - migration rate
    s - selection coefficient: residents have a viability advantage of s over 
        migrants, equivalent to migrants having a viability cost of 
        s^ = s/(1+s)"""
    return m * rvDS(s)

def gffDS(m,s):
    """Gene flow factor (gff) for the case of divergent selection (DS).
    Derived from fitness graph for the reproductive value of a migrant.
    m - migration rate
    s - selection coefficient: residents have a viability advantage of s over
        migrants, equivalent to migrants having a viability cost of 
        s^ = s/(1+s)"""
    return meDS(s) / m
    
# unidirectional CI:
# homogenous infected population:    
def rvFUHOMW(f,ci):
    """Reproductive value (rv) of an uninfected female migrant (FU) in a 
    homogenous Wolbachia infected population (HOMW). This applies to the case 
    of perfect Wolbachia transmission, t=1. Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI"""
    return (1-ci) / (1 - 2*f + ci)

def rvMUHOMW(f,ci):
    """Reproductive value (rv) of an uninfected male migrant (MU) in a 
    homogenous Wolbachia infected population (HOMW). This applies to the case 
    of perfect Wolbachia transmission, t=1. Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI"""
    return 1

def rvUHOMW(f,ci):
    """Reproductive value (rv) of an uninfected migrant (U) in a a homogenous 
    Wolbachia infected population (HOMW). This applies to the case of perfect 
    Wolbachia transmission, t=1. Derived from fitness graph. Averaged over the 
    two sexes.
    f  - fecundity reduction in infected females
    ci - level of CI"""
    return (1-f) / (1 - 2*f + ci)

# homogenous uninfected population:
def rvFWHOMU(f,ci,t):
    """Reproductive value (rv) of a Wolbachia infected female migrant (FW) in a 
    homogenous uninfected population (HOMU). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return (1-f)*(2-(1+ci)*t) / (2-(1-f)*t)

def rvMWHOMU(f,ci,t):
    """Reproductive value (rv) of a Wolbachia infected male migrant (MW) in a 
    homogenous uninfected population (HOMU). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return 1-ci

def rvWHOMU(f,ci,t):
    """Reproductive value (rv) of a Wolbachia infected migrant (W) in a 
    homogenous uninfected population (HOU). Averaged over the two sexes.
    Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return 1 - (f+ci) / (2-(1-f)*t)

def rvIPHOMU(f,ci,t):
    """Average reproductive value (rv) of a migrant from an infected population 
    (IP), i.e. where the Wolbachia infection is at a stable frequency below 1 
    due to imperfect transmission, in a homogenous uninfected population (HOU). 
    Averaged over sexes (F/M) and cytotypes (U/W). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    x = fix3SP(f,ci,t)
    return 1 - x*(f+ci) / (2-(1-f)*t)

# heterogenous population:
def rvFUHETCP(f,ci,t):
    """Reproductive value (rv) of an uninfected female (FU) in a heterogenous 
    host population (HET), 'per capita' (CP) reproductive value (as opposed to 
    'per class'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return (2*(-1+ci)*ci*(-1+f)*t)/(f*(f+sqrt(R(f,ci,t)))+power(ci,2)*(1+(-1+f)*t)-ci*(sqrt(R(f,ci,t))+power(f,2)*(5-6*t)*t+(4-3*sqrt(R(f,ci,t))-6*t)*t+f*(2+3*t*(-3+sqrt(R(f,ci,t))+4*t))))

def rvFUHETCL(f,ci,t):
    """Reproductive value (rv) of an uninfected female (FU) in a heterogenous 
    host population (HET), 'per class' (CL) reproductive value (as opposed to 
    'per capita'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return rvFUHETCP(f,ci,t) * (1-fix3SP(f,ci,t)) / 2.

def rvMUHETCP(f,ci,t):
    """Reproductive value (rv) of an uninfected male (MU) in a heterogenous 
    host population (HET), 'per capita' (CP) reproductive value (as opposed to 
    'per class'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return (ci*(power(ci,2)+(-1+f)*(f+sqrt(R(f,ci,t)))*(1+2*f*(-1+t)-2*t)*t+ci*(-f+sqrt(R(f,ci,t))-3*t+(7-4*f)*f*t+4*power((-1+f)*t,2))))/(power(ci,3)+power(f,2)*(f+sqrt(R(f,ci,t)))+power(ci,2)*(-f+sqrt(R(f,ci,t))-4*t+(9-5*f)*f*t+5*power((-1+f)*t,2))+ci*(5*power(f,3)*(-1+t)*t+sqrt(R(f,ci,t))*t*(-2+3*t)+f*t*(-4+5*sqrt(R(f,ci,t))+5*t-6*sqrt(R(f,ci,t))*t)+power(f,2)*(-1+t*(9-3*sqrt(R(f,ci,t))-10*t+3*sqrt(R(f,ci,t))*t))))

def rvMUHETCL(f,ci,t):
    """Reproductive value (rv) of an uninfected male (MU) in a heterogenous 
    host population (HET), 'per class' (CL) reproductive value (as opposed to 
    'per capita'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return rvMUHETCP(f,ci,t) * (1-fix3SP(f,ci,t)) / 2.

def rvFWHETCP(f,ci,t):
    """Reproductive value (rv) of a Wolbachia infected female (FW) in a 
    heterogenous host population (HET), 'per capita' (CP) reproductive value 
    (as opposed to 'per class'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return -(4*ci*(power(ci,4)*(1+2*f*(-1+t)-2*t)+(-1+f)*f*(f+sqrt(R(f,ci,t)))*(1+2*f*(-1+t)-2*t)*(power(f,2)*(-1+t)-t)*t+power(ci,3)*(sqrt(R(f,ci,t))-2*(3+sqrt(R(f,ci,t)))*t+13*power(f,3)*power(-1+t,2)*t+17*power(t,2)-13*power(t,3)+power(f,2)*(4+t*(-34+69*t-39*power(t,2)))+f*(-2*(1+sqrt(R(f,ci,t)))+t*(27+2*sqrt(R(f,ci,t))-60*t+39*power(t,2))))+ci*(-1+f)*t*(10*power(f,5)*power(-1+t,3)*t+power(f,4)*power(-1+t,2)*(-1+3*t*(9-2*sqrt(R(f,ci,t))+2*(-6+sqrt(R(f,ci,t)))*t))+t*(-3*sqrt(R(f,ci,t))+t*(2+9*sqrt(R(f,ci,t))+t*(-6-11*sqrt(R(f,ci,t))+4*t+6*sqrt(R(f,ci,t))*t)))-f*t*(2-11*sqrt(R(f,ci,t))+t*(1+39*sqrt(R(f,ci,t))+3*t*(-3-17*sqrt(R(f,ci,t))+2*t+8*sqrt(R(f,ci,t))*t)))-power(f,3)*(-1+t)*(sqrt(R(f,ci,t))+t*(4*(-7+4*sqrt(R(f,ci,t)))+t*(67-41*sqrt(R(f,ci,t))-44*t+24*sqrt(R(f,ci,t))*t)))+power(f,2)*(-sqrt(R(f,ci,t))+t*(11-17*sqrt(R(f,ci,t))+t*(-35+69*sqrt(R(f,ci,t))+t*(39-87*sqrt(R(f,ci,t))+4*(-4+9*sqrt(R(f,ci,t)))*t)))))+power(ci,2)*(20*power(f,5)*power(-1+t,3)*power(t,2)-power(f,4)*power(-1+t,2)*t*(8+t*(-81+100*t))+t*(-4*sqrt(R(f,ci,t))+t*(9+11*sqrt(R(f,ci,t))-t*(31+9*sqrt(R(f,ci,t))+t*(-41+20*t))))+power(f,3)*(-1+t)*(2+t*(-9*(2+sqrt(R(f,ci,t)))+t*(157+9*sqrt(R(f,ci,t))+4*t*(-81+50*t))))+power(f,2)*(1+2*sqrt(R(f,ci,t))-t*(16+22*sqrt(R(f,ci,t))+t*(-148-47*sqrt(R(f,ci,t))+t*(423+27*sqrt(R(f,ci,t))-486*t+200*power(t,2)))))+f*(-sqrt(R(f,ci,t))+t*(4+17*sqrt(R(f,ci,t))+t*(-59-40*sqrt(R(f,ci,t))+t*(3*(61+9*sqrt(R(f,ci,t)))+4*t*(-56+25*t))))))))/((ci+f+sqrt(R(f,ci,t)))*(ci-f+sqrt(R(f,ci,t))-2*t+2*(3-2*f)*f*t+4*power((-1+f)*t,2))*(power(ci,3)+power(f,2)*(f+sqrt(R(f,ci,t)))+power(ci,2)*(-f+sqrt(R(f,ci,t))-4*t+(9-5*f)*f*t+5*power((-1+f)*t,2))+ci*(5*power(f,3)*(-1+t)*t+sqrt(R(f,ci,t))*t*(-2+3*t)+f*t*(-4+5*sqrt(R(f,ci,t))+5*t-6*sqrt(R(f,ci,t))*t)+power(f,2)*(-1+t*(9-3*sqrt(R(f,ci,t))-10*t+3*sqrt(R(f,ci,t))*t)))))

def rvFWHETCL(f,ci,t):
    """Reproductive value (rv) of an infected female (FW) in a heterogenous 
    host population (HET), 'per class' (CL) reproductive value (as opposed to 
    'per capita'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return rvFWHETCP(f,ci,t) * fix3SP(f,ci,t) / 2.

def rvMWHETCP(f,ci,t):
    """Reproductive value (rv) of a Wolbachia infected male (MW) in a 
    heterogenous host population (HET), 'per capita' (CP) reproductive value 
    (as opposed to 'per class'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return (4*ci*(power(ci,4)+(-1+f)*f*(f+sqrt(R(f,ci,t)))*(1+2*f*(-1+t)-2*t)*t*(power(f,2)+(-1+(4-3*f)*f)*t+2*power((-1+f)*t,2))+power(ci,3)*(-2*f+sqrt(R(f,ci,t))-6*t-f*(-13+f*(6+f))*t+power(-1+f,2)*(7+2*f)*power(t,2)-power(-1+f,3)*power(t,3))+power(ci,2)*(-4*power(f,5)*power(-1+t,3)*power(t,2)+power(f,4)*power(-1+t,2)*power(t,2)*(3+20*t)-power(f,3)*(-1+t)*t*(6-sqrt(R(f,ci,t))+t*(-41+sqrt(R(f,ci,t))+4*t*(3+10*t)))+t*(-4*sqrt(R(f,ci,t))+t*(9+5*sqrt(R(f,ci,t))+t*(-23+sqrt(R(f,ci,t))+t*(11+4*t))))-f*(sqrt(R(f,ci,t))+t*(-4-9*sqrt(R(f,ci,t))+t*(47+8*sqrt(R(f,ci,t))+t*(-91+3*sqrt(R(f,ci,t))+4*t*(8+5*t)))))+power(f,2)*(1+t*(-2*(5+2*sqrt(R(f,ci,t)))+t*(78+sqrt(R(f,ci,t))+t*(3*(-41+sqrt(R(f,ci,t)))+2*t*(9+20*t))))))+ci*(-1+f)*t*(2*power(f,5)*power(-1+t,3)*t*(-3+4*t)-t*(-1+2*t)*(-3*sqrt(R(f,ci,t))+t*(2+3*sqrt(R(f,ci,t))+t*(-6+sqrt(R(f,ci,t))+4*t)))-power(f,4)*power(-1+t,2)*(-1+t*(11-2*sqrt(R(f,ci,t))+2*t*(-28+sqrt(R(f,ci,t))+20*t)))+power(f,3)*(-1+t)*(sqrt(R(f,ci,t))+t*(-4*(1+sqrt(R(f,ci,t)))+t*(83-5*sqrt(R(f,ci,t))+4*t*(-41+2*sqrt(R(f,ci,t))+20*t))))+f*t*(-2+13*sqrt(R(f,ci,t))+t*(-11-29*sqrt(R(f,ci,t))+t*(67+9*sqrt(R(f,ci,t))+2*t*(-47+4*sqrt(R(f,ci,t))+20*t))))+power(f,2)*(sqrt(R(f,ci,t))+t*(5-17*sqrt(R(f,ci,t))+t*(43+25*sqrt(R(f,ci,t))+t*(3*(-61+sqrt(R(f,ci,t)))-4*t*(-54+3*sqrt(R(f,ci,t))+20*t))))))))/((ci+f+sqrt(R(f,ci,t)))*(ci-f+sqrt(R(f,ci,t))-2*t+2*(3-2*f)*f*t+4*power((-1+f)*t,2))*(power(ci,3)+power(f,2)*(f+sqrt(R(f,ci,t)))+power(ci,2)*(-f+sqrt(R(f,ci,t))-4*t+(9-5*f)*f*t+5*power((-1+f)*t,2))+ci*(5*power(f,3)*(-1+t)*t+sqrt(R(f,ci,t))*t*(-2+3*t)+f*t*(-4+5*sqrt(R(f,ci,t))+5*t-6*sqrt(R(f,ci,t))*t)+power(f,2)*(-1+t*(9-3*sqrt(R(f,ci,t))-10*t+3*sqrt(R(f,ci,t))*t)))))

def rvMWHETCL(f,ci,t):
    """Reproductive value (rv) of an infected male (MW) in a heterogenous host 
    population (HET), 'per class' (CL) reproductive value (as opposed to 'per 
    capita'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return rvMWHETCP(f,ci,t) * fix3SP(f,ci,t) / 2.

def rvUHETCP(f,ci,t):
    """Average reproductive value (rv) of an uninfected host (U) in a 
    heterogenous population (HET), 'per capita' (CP) reproductive value (as 
    opposed to 'per class'). Derived from fitness graph. 
    Averaged over the two sexes.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return ((-1+f)*t*(power(ci,2)+f*(f-sqrt(R(f,ci,t)))-ci*(sqrt(R(f,ci,t))+4*t+2*(f+f*(-5+3*f)*t-3*power((-1+f)*t,2)))))/(2*(power(ci-f,2)+2*(-1+f)*(ci*(2+ci)-4*ci*f+power(f,2))*t-3*ci*power(-1+f,2)*(-4+3*f)*power(t,2)+9*ci*power((-1+f)*t,3)))

def rvUHETCL(f,ci,t):
    """Average reproductive value (rv) of an uninfected host (U) in a 
    heterogenous population (HET), 'per class' (CL) reproductive value (as 
    opposed to 'per capita'). Derived from fitness graph. 
    Averaged over the two sexes.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia"""
    return rvUHETCP(f,ci,t) * (1-fix3SP(f,ci,t))

def rvWHETCP(f,ci,t):
    """
    Average reproductive value (rv) of a Wolbachia infected host (W) in a 
    heterogenous population (HET), 'per capita' (CP) reproductive value (as 
    opposed to 'per class'). Derived from fitness graph. Averaged over the two sexes.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia
    """
    return (4*ci*(f*(-1+t)-t)*(-power(ci,4)+(-1+f)*f*(f+sqrt(R(f,ci,t)))*(1+2*f*(-1+t)-2*t)*(f*(-1+t)-t)*t-power(ci,3)*(-2*f+sqrt(R(f,ci,t))-5*t+(12-7*f)*f*t+7*power((-1+f)*t,2))-ci*(-1+f)*t*(-4*power(f,4)*(-2+t)*power(-1+t,2)*t+power(t,2)*(-2+3*sqrt(R(f,ci,t))-2*t*(-3+2*sqrt(R(f,ci,t))+2*t))+power(f,3)*(-1+t)*(-1+t*(19-4*sqrt(R(f,ci,t))+2*t*(-19+2*sqrt(R(f,ci,t))+8*t)))+f*(sqrt(R(f,ci,t))+(-1+t)*t*(3-2*sqrt(R(f,ci,t))+2*t*(-9+6*sqrt(R(f,ci,t))+8*t)))-power(f,2)*(sqrt(R(f,ci,t))+t*(-15+6*sqrt(R(f,ci,t))+t*(56-19*sqrt(R(f,ci,t))+6*t*(-11+2*sqrt(R(f,ci,t))+4*t)))))-power(ci,2)*(12*power(f,4)*power((-1+t)*t,2)+power(f,3)*t*(4+t*(-43+87*t-48*power(t,2)))+t*(-3*sqrt(R(f,ci,t))+t*(4+5*sqrt(R(f,ci,t))+3*t*(-5+4*t)))-f*(sqrt(R(f,ci,t))+t*(-2-8*sqrt(R(f,ci,t))+t*(27+10*sqrt(R(f,ci,t))-69*t+48*power(t,2))))+power(f,2)*(1+t*(-6-5*sqrt(R(f,ci,t))+t*(54+5*sqrt(R(f,ci,t))+9*t*(-13+8*t)))))))/((ci+f+sqrt(R(f,ci,t)))*(ci-f+sqrt(R(f,ci,t))-2*t+2*(3-2*f)*f*t+4*power((-1+f)*t,2))*(power(ci,3)+power(f,2)*(f+sqrt(R(f,ci,t)))+power(ci,2)*(-f+sqrt(R(f,ci,t))-4*t+(9-5*f)*f*t+5*power((-1+f)*t,2))+ci*(5*power(f,3)*(-1+t)*t+sqrt(R(f,ci,t))*t*(-2+3*t)+f*t*(-4+5*sqrt(R(f,ci,t))+5*t-6*sqrt(R(f,ci,t))*t)+power(f,2)*(-1+t*(9-3*sqrt(R(f,ci,t))-10*t+3*sqrt(R(f,ci,t))*t)))))

def rvWHETCL(f,ci,t):
    """
    Reproductive value (rv) of an infected host (W) in a heterogenous host 
    population (HET), 'per class' (CL) reproductive value (as opposed to 'per 
    capita'). Derived from fitness graph.
    f  - fecundity reduction in infected females
    ci - level of CI
    t  - transmission rate of Wolbachia
    """
    return rvWHETCP(f,ci,t) * fix3SP(f,ci,t)

################################## Reinforcement ##############################################
def xPref(ci,pr,q,s):
    """
    Approximated equilibrium frequency of a mutant allele at the locus for 
    female mating preference in an uninfected island receiving migration from 
    an infected mainland.
    ci - level of CI
    pr - rejection probability of mating preference mutant allele
    q  - transition probability to a new mating round
    s  - selection coefficient (the locally adaptive trait has a 1+s times higher fitness than the 
         mainland trait)
    """
    return 1 - (1-ci)*s/((2*ci*s - (1-ci)*(1-q))*pr)

def prcrit(ci,q,s):
    """
    Minimal rejection probability (pr) of a mutant allele at the locus for 
    female mating preference to spread in an uninfected island receiving 
    migration from an infected mainland.
    ci - level of CI
    q  - transition probability to a new mating round
    s  - selection coefficient (the locally adaptive trait has a 1+s times higher fitness than the 
         mainland trait)
    """
    return (1-ci)*s / (2*ci*s - (1-ci)*(1-q))

def xTprefEQ(pr,s,xPref):
    """Modified from Kirkpatrick (1982), equation (2).
    """
    # calculate Kirkpatrick's preference strength (ak) from rejection 
    # probability (pr):
    ak = 1./(1-pr)    
    # calculate Kirkpatrick's viability cost (sk) from selection 
    # coefficient (s):
    sk = s/(1.+s)
    # Kirkpatrick used a male-limited trait, but in our model, the trait is 
    # expressed in both sexes     :
    p2 = xPref/2.
    
    if p2 <= sk / ((ak-1)*(1-sk)):
        return 0.
    elif p2 >= ak*sk / (ak-1):
        return 1.
    else:
        return (1./sk + 1./(ak*(1-sk)-1))*p2 - 1./(ak*(1-sk)-1)
    
    
